#pragma once

#ifndef NTFSUTILS_H
#define NTFSUTILS_H

#include "ntfs.h"
#include "disk.h"
#include "sxmem.h"
#include <vector>

using namespace std;
using namespace sx;

namespace ntfs 
{

/******************************************************************************/
/*                                  CNTFSRun                                  */
/******************************************************************************/

class CNonResidentAttribute;

typedef struct tagRUNITEM 
{
    UINT64 LowVCN;
    UINT64 HighVCN;
    UINT64 LowLCN;
    UINT64 HighLCN;
    UINT64 ClustersCount;
    tagRUNITEM(): LowVCN(0), HighVCN(0), LowLCN(0), HighLCN(0), ClustersCount(0) {}
} RUNITEM, *PRUNITEM;

class CNTFSRun
{
public:              
    CNTFSRun(const PCNONRESIDENT_ATTRIBUTE pcNonResidentAttribute);
    UINT64 LCN(UINT64 VCN) const;
    UINT64 LowVCN;
    UINT64 HighVCN;
private:
    vector<RUNITEM> mRunList;
    friend void dumpAttributeNonResident(CNonResidentAttribute* pAttr);
};

typedef CNTFSRun CRun;

/******************************************************************************/
/*                                  CAttribure                                */
/******************************************************************************/

class CAttribute
{
public:
    CAttribute(const SxSharedBuffer& Buffer): mBuffer(Buffer) {}
    virtual ~CAttribute() {}

    PCATTRIBUTE operator->() const { return mBuffer; }

    operator PATTRIBUTE()  { return mBuffer; }
    operator PCATTRIBUTE() const  { return mBuffer; }
protected:
    SxSharedBuffer mBuffer;
};

class CResidentAttribute: public CAttribute
{
public:
    CResidentAttribute(const SxSharedBuffer mBuffer)
        :CAttribute(mBuffer)
    {}

    PCRESIDENT_ATTRIBUTE operator->() const { return mBuffer; }

    operator PRESIDENT_ATTRIBUTE()  { return mBuffer; }
    operator PCRESIDENT_ATTRIBUTE() const  { return mBuffer; }
    
    friend SxSharedBuffer getAttributeValue(const CResidentAttribute& Attribute);
};

class CNTFSVolume;

class CNonResidentAttribute: public CAttribute
{
public:
    CNonResidentAttribute(const SxSharedBuffer mBuffer)
        :CAttribute(mBuffer), mRun(mBuffer)
    {}

    PCNONRESIDENT_ATTRIBUTE operator->() const { return mBuffer; }

    operator PNONRESIDENT_ATTRIBUTE()  { return mBuffer; }
    operator PCNONRESIDENT_ATTRIBUTE() const  { return mBuffer; }

    UINT64 getLCN(UINT64 VCN) const { return mRun.LCN(VCN); }
    UINT64 getLowVCN() const { return mRun.LowVCN; }
    UINT64 getHighVCN() const { return mRun.HighVCN; }
    
private:
    CRun mRun;
    friend void dumpAttributeNonResident(CNonResidentAttribute* pAttr);
    friend SxSharedBuffer getAttributeValue(const CNonResidentAttribute& Attribute, const CNTFSVolume& Volume);
};

SxSharedBuffer getAttributeValue(const CResidentAttribute& Attribute);
SxSharedBuffer getAttributeValue(const CNonResidentAttribute& Attribute, const CNTFSVolume& Volume);

/******************************************************************************/
/*                                CNTFSVolume                                 */
/******************************************************************************/

class CNTFSVolume: public CDisk
{
public:
    CNTFSVolume(const tstring fileName);
    ~CNTFSVolume();
    BOOT_BLOCK getBootBlock() const { return mBootBlock; }
    UINT64 getMFTOffset() const { return mMFTOffset; }
    UINT32 getFileRecordLength() const { return mBytesPerFileRecord; }
    UINT32 getClusterLength() const { return mBytesPerCluster; }
    SxSharedBuffer readCluster(UINT64 LCN) const;
    
    BOOT_BLOCK mBootBlock;
private:

    UINT64 mMFTOffset;
    UINT32 mBytesPerFileRecord;
    UINT32 mBytesPerCluster;
};

/******************************************************************************/
/*                                CFileRecord                                 */
/******************************************************************************/

class CFileRecord 
{
public:
    CFileRecord(const SxSharedBuffer& MemBuffer);
    CFileRecord(): mBuffer(0) {}
    CFileRecord(const CFileRecord& FileRecord);
    CFileRecord& operator =(const CFileRecord& FileRecord);
    virtual ~CFileRecord();

    CAttribute* getAttribute(unsigned int Index) 
    {
        return const_cast<CAttribute*>
            (static_cast<const CFileRecord&>(*this).getAttribute(Index));
    }
    const CAttribute* getAttribute(unsigned int Index) const { return mAttributeList[Index]; }
    unsigned int getAttributesCount() const { return mAttributeList.size(); }

    CAttribute* getAttribute(ATTRIBUTE_TYPE AttributeType) 
    { 
        return const_cast<CAttribute*>
            (static_cast<const CFileRecord*>(this)->getAttribute(AttributeType)); 
    }
    const CAttribute* getAttribute(ATTRIBUTE_TYPE AttributeType) const ;

    PCFILE_RECORD_HEADER operator->() const { return mBuffer; }

    SxSharedBuffer mBuffer;
protected:
    void LoadAttributes();

private:
    vector<CAttribute*> mAttributeList;
};

/******************************************************************************/
/*                                    CMFT                                    */
/******************************************************************************/

class CMFT : public CFileRecord 
{
public:
    CMFT(CNTFSVolume* pVolume);
    CFileRecord getFileRecord(UINT64 Index);
private:
    CNTFSVolume* mpVolume;
};

}  //  namespace ntfs

#endif  //  NTFSUTILS_H
